import yaml
import networkx as nx
import matplotlib.pyplot as plt
import sys
import os

def load_topology_from_yaml(file_path):
    """加载 YAML 文件并返回节点数据和格式类型"""
    if not os.path.exists(file_path):
        print(f"错误：文件 '{file_path}' 不存在！")
        sys.exit(1)
    
    with open(file_path, 'r', encoding='utf-8') as f:
        try:
            data = yaml.safe_load(f)
        except yaml.YAMLError as e:
            print(f"YAML 解析错误: {e}")
            sys.exit(1)
    
    if not isinstance(data, dict):
        print("错误：YAML 文件应为字典格式")
        sys.exit(1)
    
    if 'node_list' not in data:
        print("错误：YAML 文件应包含顶层键 'node_list'")
        sys.exit(1)
    
    nodes = data['node_list']
    
    if not isinstance(nodes, list):
        print(f"错误：节点数据应为列表，当前类型: {type(nodes)}")
        sys.exit(1)
    
    # 检测格式类型：通过检查第一个节点的 param 内容
    if nodes and len(nodes) > 0:
        first_param = nodes[0].get('param', [])
        if isinstance(first_param, list) and len(first_param) >= 6:
            # stairs 格式：param 包含 [x, y, z, yaw, map_name, parent, type]
            format_type = 'stairs'
        else:
            # map 格式：param 包含邻居节点列表
            format_type = 'map'
    else:
        format_type = 'map'
    
    return nodes, format_type

def build_graph_stairs(nodes):
    """构建 stairs 格式的图（包含坐标信息）"""
    G = nx.Graph()
    positions = {}
    
    for node in nodes:
        if not isinstance(node, dict):
            print(f"警告：跳过非字典项: {node}")
            continue
        if 'node_name' not in node:
            print(f"警告：节点缺少 'node_name' 字段，跳过: {node}")
            continue
        
        name = node['node_name']
        params = node.get('param', [])
        
        if not isinstance(params, list) or len(params) < 6:
            print(f"警告：节点 '{name}' 的 'param' 格式不正确，跳过")
            continue
        
        # 解析参数：[x, y, z, yaw, map_name, parent_node, type]
        x, y = float(params[0]), float(params[1])
        parent_node = params[5]
        
        # 添加节点并保存位置
        G.add_node(name)
        positions[name] = (x, y)
        
        # 如果有父节点，添加边
        if parent_node and parent_node != '':
            G.add_edge(name, parent_node)
    
    return G, positions

def build_graph_map(nodes):
    """构建 map 格式的图（拓扑连接）"""
    G = nx.Graph()
    
    for node in nodes:
        if not isinstance(node, dict):
            print(f"警告：跳过非字典项: {node}")
            continue
        
        # 使用 node_name 字段
        name = node.get('node_name')
        if not name:
            print(f"警告：节点缺少 'node_name' 字段，跳过: {node}")
            continue
        
        G.add_node(name)
        params = node.get('param', [])
        
        if not isinstance(params, list):
            print(f"警告：节点 '{name}' 的 'param' 不是列表，跳过连接")
            continue
        
        # 添加所有邻居边
        for neighbor in params:
            if isinstance(neighbor, str):
                G.add_edge(name, neighbor)
    
    return G, None

def visualize_graph(G, positions=None, title="Topological Road Network"):
    """可视化图"""
    plt.figure(figsize=(12, 10))
    
    # 如果有位置信息，使用实际坐标；否则使用自动布局
    if positions:
        pos = positions
    else:
        pos = nx.spring_layout(G, seed=42, k=2)
    
    # 绘制节点
    nx.draw_networkx_nodes(G, pos, node_size=2500, node_color='lightcoral', 
                          edgecolors='black', linewidths=2)
    
    # 绘制边
    nx.draw_networkx_edges(G, pos, width=2.5, alpha=0.6, edge_color='gray')
    
    # 绘制标签
    nx.draw_networkx_labels(G, pos, font_size=11, font_weight='bold')
    
    # 如果有坐标信息，添加坐标标注
    if positions:
        coord_labels = {node: f"({positions[node][0]:.2f}, {positions[node][1]:.2f})" 
                       for node in positions}
        nx.draw_networkx_labels(G, pos, coord_labels, font_size=8, 
                              font_color='blue', verticalalignment='bottom')
        plt.axis('equal')
        plt.grid(True, alpha=0.3)
    else:
        plt.axis('off')
    
    plt.title(title, fontsize=16, fontweight='bold')
    plt.tight_layout()
    plt.show()

def main():
    # 支持命令行参数
    if len(sys.argv) > 1:
        yaml_file = sys.argv[1]
    else:
        yaml_file = "map_topo.yaml"
    
    print(f"正在从 '{yaml_file}' 加载拓扑数据...")
    nodes, format_type = load_topology_from_yaml(yaml_file)
    
    print(f"检测到格式类型: {format_type}")
    
    # 根据格式构建图
    if format_type == 'stairs':
        G, positions = build_graph_stairs(nodes)
    else:
        G, positions = build_graph_map(nodes)
    
    print(f"✅ 图构建完成：{G.number_of_nodes()} 个节点，{G.number_of_edges()} 条边")
    
    # 显示节点信息
    print("\n节点列表:")
    for node in G.nodes():
        neighbors = list(G.neighbors(node))
        if positions and node in positions:
            print(f"  - {node}: 位置 {positions[node]}, 连接到 {neighbors}")
        else:
            print(f"  - {node}: 连接到 {neighbors}")
    
    # 可视化
    title = f"Topological Network - {os.path.basename(yaml_file)}"
    visualize_graph(G, positions, title)

if __name__ == "__main__":
    main()
